package com.snap.vseries.analyze;

import org.apache.commons.math3.analysis.ParametricUnivariateFunction;
import org.apache.commons.math3.fitting.AbstractCurveFitter;
import org.apache.commons.math3.fitting.WeightedObservedPoint;
import org.apache.commons.math3.fitting.leastsquares.LeastSquaresBuilder;
import org.apache.commons.math3.fitting.leastsquares.LeastSquaresProblem;
import org.apache.commons.math3.fitting.leastsquares.ParameterValidator;
import org.apache.commons.math3.linear.DiagonalMatrix;

import java.util.Collection;

/**
 * 曲线拟合器，使用特性：
 * 1.懒评估lazyEvaluation
 * 2.参数校验器parameterValidator
 */
public class MyCurveFitter extends AbstractCurveFitter {

    private static final String TAG = "MyCurveFitter";

    private final ParametricUnivariateFunction function;
    private final double[] initialGuess;
    private final int maxIter;
    private final ParameterValidator validator;

    public MyCurveFitter(ParametricUnivariateFunction function, double[] initialGuess, int maxIter, ParameterValidator validator) {
        this.function = function;
        this.initialGuess = initialGuess;
        this.maxIter = maxIter;
        this.validator = validator;
    }

    @Override
    protected LeastSquaresProblem getProblem(Collection<WeightedObservedPoint> observations) {
        // Prepare least-squares problem.
        final int len = observations.size();
        final double[] target = new double[len];
        final double[] weights = new double[len];

        int count = 0;
        for (WeightedObservedPoint obs : observations) {
            target[count] = obs.getY();
            weights[count] = obs.getWeight();
            ++count;
        }

        final AbstractCurveFitter.TheoreticalValuesFunction model
                = new AbstractCurveFitter.TheoreticalValuesFunction(function, observations);

        // Create an optimizer for fitting the curve to the observed points.
        return new LeastSquaresBuilder().
                maxEvaluations(Integer.MAX_VALUE).
                maxIterations(maxIter).
                start(initialGuess).
                target(target).
                weight(new DiagonalMatrix(weights)).
                model(model.getModelFunction(), model.getModelFunctionJacobian()).
                lazyEvaluation(true).
                parameterValidator(validator).
                build();
    }
}
